#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文档重置脚本
清除数据库中的文档记录，重新扫描和处理所有文档
"""

import sqlite3
import os
import sys
from pathlib import Path

def reset_documents_db(db_path: str = "data/database/tiny_qa.db"):
    """重置文档数据库"""
    try:
        with sqlite3.connect(db_path) as conn:
            cursor = conn.cursor()
            
            print("🗑️ 清除文档数据...")
            
            # 删除文档块
            cursor.execute("DELETE FROM document_chunks")
            print(f"   - 删除了 {cursor.rowcount} 个文档块")
            
            # 删除文档记录
            cursor.execute("DELETE FROM documents")
            print(f"   - 删除了 {cursor.rowcount} 个文档记录")
            
            # 重置自增ID
            cursor.execute("DELETE FROM sqlite_sequence WHERE name IN ('documents', 'document_chunks')")
            
            conn.commit()
            print("✅ 数据库重置完成")
            
    except Exception as e:
        print(f"❌ 重置数据库失败: {e}")
        return False
    
    return True

def clear_vector_files(vector_path: str = "data/vectors"):
    """清除向量文件"""
    try:
        if os.path.exists(vector_path):
            vector_files = [f for f in os.listdir(vector_path) if f.endswith('.json')]
            for file in vector_files:
                os.remove(os.path.join(vector_path, file))
            print(f"🗑️ 删除了 {len(vector_files)} 个向量文件")
        else:
            print("ℹ️ 向量目录不存在，跳过")
    except Exception as e:
        print(f"⚠️ 清除向量文件失败: {e}")

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='重置文档数据库')
    parser.add_argument('--confirm', action='store_true', help='确认重置操作')
    parser.add_argument('--db', default='data/database/tiny_qa.db', help='数据库路径')
    parser.add_argument('--vectors', default='data/vectors', help='向量目录路径')
    
    args = parser.parse_args()
    
    if not args.confirm:
        print("⚠️ 警告：此操作将删除所有文档数据和向量文件！")
        print("使用 --confirm 参数确认执行")
        sys.exit(1)
    
    print("🔄 开始重置文档系统...")
    print("=" * 50)
    
    # 1. 清除向量文件
    clear_vector_files(args.vectors)
    
    # 2. 重置数据库
    if reset_documents_db(args.db):
        print("\n🎉 重置完成！")
        print("\n📝 下一步操作：")
        print("1. 访问 http://localhost:5002/documents")
        print("2. 点击'扫描文档'发现文档")
        print("3. 点击'处理所有文档'重新处理")
        print("\n或者运行测试脚本：")
        print("python3 scripts/test_qa_system.py")
    else:
        print("\n❌ 重置失败")
        sys.exit(1)

if __name__ == "__main__":
    main() 